import numpy as np
import os
from os import path

def relu(x):
    ret = np.zeros_like(x)
    ret[x>0] = x[x>0]
    return ret

def F_to_Q4_12(f):
    return np.round(f * (1<<12)).astype(np.int16)

def Q4_12_to_F(q):
    return (q / (1<<12)).astype(np.float32)

def np2str(a: np.array)-> str:
    cnt, s = 0, ''
    for var in a.flatten():
        s += str(var) + ', '
        if cnt %8 == 7:
            s += '\n'
        cnt += 1
    return s


#################
def save_np_as_q4_12_c_header1(a, fname):
    fmt = '''
#ifndef {0}
#define {0}

const signed short int {1}[{2}] = {{
{3}}};

#endif
'''
    if path.exists(fname+'.h'):
        os.remove(fname+'.h')
        
    with open(fname+'.h', 'w') as fp:
        fp.write(fmt.format(
            '_' + fname.upper() + '__H_',
            fname,
            a.shape[0],
            np2str(F_to_Q4_12(a))
        ))

#################
def save_np_as_q4_12_c_header2(a, fname):
    fmt = '''
#ifndef {0}
#define {0}

const signed short int {1}[{2}][{3}] = {{
{4}}};

#endif
'''
    if path.exists(fname+'.h'):
        os.remove(fname+'.h')
        
    with open(fname+'.h', 'w') as fp:
        fp.write(fmt.format(
            '_' + fname.upper() + '__H_',
            fname,
            a.shape[0],
            a.shape[1],
            np2str(F_to_Q4_12(a))
        ))

np.random.seed(1)

x = np.random.randn(84,1) 
w = np.random.randn(10,84) / 10
b = np.random.randn(10,1)

y = w@x+b
y = relu(w@x+b)
save_np_as_q4_12_c_header1(x, "fc3_din")
save_np_as_q4_12_c_header2(w, "fc3_w")
save_np_as_q4_12_c_header1(b, "fc3_b")
save_np_as_q4_12_c_header1(y, "fc3_dout")

print(F_to_Q4_12(y))